import argparse
import os
from face_searcher import FaceSearcher
import logging

def setup_logger():
    """配置日志"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    return logging.getLogger(__name__)

def add_faces(searcher: FaceSearcher, image_path: str, person_id: str, metadata: dict = None):
    """添加单个人脸或目录中的所有人脸"""
    if os.path.isdir(image_path):
        # 批量添加目录中的人脸
        count = searcher.batch_add_faces(image_path, person_id)
        print(f"成功添加 {count} 张人脸照片到数据库")
    else:
        # 添加单张人脸照片
        success = searcher.add_face(image_path, person_id, metadata)
        if success:
            print(f"成功添加人脸照片: {image_path}")
        else:
            print(f"添加人脸失败: {image_path}")

def search_faces(searcher: FaceSearcher, query_path: str, threshold: float = 0.6, top_k: int = 5):
    """搜索相似人脸"""
    if os.path.isdir(query_path):
        # 批量搜索目录中的所有图片
        image_paths = []
        for filename in os.listdir(query_path):
            if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
                image_paths.append(os.path.join(query_path, filename))
        
        results = searcher.batch_search(image_paths, threshold, top_k)
        
        # 打印结果
        for query_image, matches in results.items():
            print(f"\n查询图片: {query_image}")
            if matches:
                print(f"找到 {len(matches)} 个相似人脸:")
                for match in matches:
                    print(f"  人物ID: {match['person_id']}")
                    print(f"  相似度: {match['similarity']:.2f}")
                    print(f"  图片路径: {match['image_path']}")
                    if match['metadata']:
                        print(f"  元数据: {match['metadata']}")
            else:
                print("未找到相似人脸")
    else:
        # 搜索单张图片
        results = searcher.search_similar_faces(query_path, threshold, top_k)
        
        if results:
            print(f"\n找到 {len(results)} 个相似人脸:")
            for result in results:
                print(f"  人物ID: {result['person_id']}")
                print(f"  相似度: {result['similarity']:.2f}")
                print(f"  图片路径: {result['image_path']}")
                if result['metadata']:
                    print(f"  元数据: {result['metadata']}")
        else:
            print("\n未找到相似人脸")

def main():
    parser = argparse.ArgumentParser(description='人脸识别与搜索工具')
    
    # 添加子命令
    subparsers = parser.add_subparsers(dest='command', help='可用命令')
    
    # 添加人脸的命令
    add_parser = subparsers.add_parser('add', help='添加人脸到数据库')
    add_parser.add_argument('image_path', help='人脸图片或包含人脸图片的目录路径')
    add_parser.add_argument('person_id', help='人物ID')
    add_parser.add_argument('--name', help='人物姓名')
    add_parser.add_argument('--age', type=int, help='人物年龄')
    add_parser.add_argument('--gender', choices=['male', 'female'], help='性别')
    
    # 搜索人脸的命令
    search_parser = subparsers.add_parser('search', help='搜索相似人脸')
    search_parser.add_argument('query_path', help='查询图片或包含查询图片的目录路径')
    search_parser.add_argument('--threshold', type=float, default=0.6, help='相似度阈值 (0-1)')
    search_parser.add_argument('--top-k', type=int, default=5, help='返回的最相似人脸数量')
    
    # MongoDB 连接参数
    parser.add_argument('--mongo-uri', 
                       default='mongodb://192.168.0.234:27017/',  # 更新为您的MongoDB服务器地址
                       help='MongoDB连接URI')
    
    args = parser.parse_args()
    
    # 初始化日志
    logger = setup_logger()
    
    try:
        # 初始化人脸搜索器
        searcher = FaceSearcher(mongo_uri=args.mongo_uri)
        
        if args.command == 'add':
            # 构建元数据
            metadata = {}
            if args.name:
                metadata['name'] = args.name
            if args.age:
                metadata['age'] = args.age
            if args.gender:
                metadata['gender'] = args.gender
                
            add_faces(searcher, args.image_path, args.person_id, metadata)
            
        elif args.command == 'search':
            search_faces(searcher, args.query_path, args.threshold, args.top_k)
            
        else:
            parser.print_help()
            
    except Exception as e:
        logger.error(f"发生错误: {str(e)}")
        raise

if __name__ == "__main__":
    main() 